#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
ROS2-NS3 API 服务器
提供前端与ROS2和NS-3仿真环境之间的通信接口
实现REST API和WebSocket，用于配置参数和实时数据交换
"""

import os
import json
import threading
import subprocess
import time
from typing import Dict, List, Any, Optional

# Flask相关导入
from flask import Flask, request, jsonify
from flask_socketio import SocketIO, emit

# ROS2相关导入
import rclpy
from rclpy.node import Node
from std_msgs.msg import Int16MultiArray, Float32MultiArray
from rclpy.qos import QoSProfile, QoSReliabilityPolicy, QoSHistoryPolicy

# 创建Flask应用
app = Flask(__name__)
# 配置密钥（在生产环境中应使用环境变量）
app.config['SECRET_KEY'] = 'ros2-ns3-secret'
# 初始化SocketIO，支持跨域请求
socketio = SocketIO(app, cors_allowed_origins="*")

# 全局变量，存储当前系统状态和配置
system_status = {
    "running": False,
    "ns3_process": None,
    "ros2_initialized": False,
    "error_message": None
}

# 默认网络配置参数
default_network_config = {
    "backbones": 8,       # 骨干UAV数量
    "dynamic": True,      # 是否为动态仿真
    "server_port": 12345, # 与ROS通信的端口
    "exp": 2.4,           # 路径损耗模型指数
    "ref_loss": 46,       # 参考路径损耗值
    "tx_power": 16.0,     # 发射功率(dBm)
    "fading_mean": 0,     # 衰落信道高斯噪声均值
    "fading_var": 32      # 衰落信道高斯噪声方差
}

# 当前网络配置，初始化为默认值
current_network_config = default_network_config.copy()

# 最新的网络数据
latest_network_data = {
    "routing_table": [],   # 存储最新的路由表
    "rss_matrix": [],      # 存储最新的RSS矩阵
    "throughput_matrix": [] # 存储最新的吞吐量矩阵
}

# ROS2 Node 类，用于与ROS2环境通信
class ROS2NS3BridgeNode(Node):
    def __init__(self):
        """初始化ROS2节点，设置订阅和发布"""
        super().__init__("ros2_ns3_bridge_node")
        
        # 设置QoS配置
        qos_profile = QoSProfile(
            reliability=QoSReliabilityPolicy.RELIABLE,
            history=QoSHistoryPolicy.KEEP_LAST,
            depth=10
        )
        
        # 创建订阅者
        self.routing_table_sub = self.create_subscription(
            Int16MultiArray,
            "/routing_table",
            self.routing_table_callback,
            qos_profile
        )
        
        self.rss_matrix_sub = self.create_subscription(
            Float32MultiArray,
            "/rosns3_client/rss_matrix",
            self.rss_matrix_callback,
            qos_profile
        )
        
        self.throughput_matrix_sub = self.create_subscription(
            Float32MultiArray,
            "/rosns3_client/throughput_matrix",
            self.throughput_matrix_callback,
            qos_profile
        )
        
        # 日志信息
        self.get_logger().info("ROS2-NS3 桥接节点已初始化")
    
    def routing_table_callback(self, msg: Int16MultiArray):
        """处理路由表数据回调"""
        # 从一维数组转换为二维邻接矩阵
        size = int(len(msg.data) ** 0.5)  # 假设是方形矩阵
        matrix = []
        for i in range(size):
            row = [msg.data[i * size + j] for j in range(size)]
            matrix.append(row)
        
        # 更新全局变量
        latest_network_data["routing_table"] = matrix
        
        # 通过WebSocket发送更新
        try:
            socketio.emit('network_update', {'type': 'routing_table', 'data': matrix})
        except Exception as e:
            self.get_logger().error(f"WebSocket发送路由表失败: {str(e)}")
        
        self.get_logger().debug("接收到新的路由表数据")
    
    def rss_matrix_callback(self, msg: Float32MultiArray):
        """处理RSS矩阵数据回调"""
        # 从一维数组转换为二维矩阵
        size = int(len(msg.data) ** 0.5)  # 假设是方形矩阵
        matrix = []
        for i in range(size):
            row = [msg.data[i * size + j] for j in range(size)]
            matrix.append(row)
        
        # 更新全局变量
        latest_network_data["rss_matrix"] = matrix
        
        # 通过WebSocket发送更新
        try:
            socketio.emit('network_update', {'type': 'rss_matrix', 'data': matrix})
        except Exception as e:
            self.get_logger().error(f"WebSocket发送RSS矩阵失败: {str(e)}")
        
        self.get_logger().debug("接收到新的RSS矩阵数据")
    
    def throughput_matrix_callback(self, msg: Float32MultiArray):
        """处理吞吐量矩阵数据回调"""
        # 从一维数组转换为二维矩阵
        size = int(len(msg.data) ** 0.5)  # 假设是方形矩阵
        matrix = []
        for i in range(size):
            row = [msg.data[i * size + j] for j in range(size)]
            matrix.append(row)
        
        # 更新全局变量
        latest_network_data["throughput_matrix"] = matrix
        
        # 通过WebSocket发送更新
        try:
            socketio.emit('network_update', {'type': 'throughput_matrix', 'data': matrix})
        except Exception as e:
            self.get_logger().error(f"WebSocket发送吞吐量矩阵失败: {str(e)}")
        
        self.get_logger().debug("接收到新的吞吐量矩阵数据")

# ROS2 节点运行线程
def run_ros2_node():
    """在独立线程中运行ROS2节点"""
    global system_status
    
    try:
        rclpy.init()
        node = ROS2NS3BridgeNode()
        system_status["ros2_initialized"] = True
        
        # 循环处理ROS2消息
        try:
            rclpy.spin(node)
        except KeyboardInterrupt:
            pass
        finally:
            # 清理资源
            node.destroy_node()
            rclpy.shutdown()
            system_status["ros2_initialized"] = False
    except Exception as e:
        system_status["error_message"] = f"ROS2节点初始化失败: {str(e)}"
        system_status["ros2_initialized"] = False

# 启动NS-3仿真
def start_ns3_simulation(config: Dict[str, Any]) -> bool:
    """
    启动NS-3仿真
    
    Args:
        config: 网络配置参数
        
    Returns:
        bool: 启动是否成功
    """
    global system_status
    
    if system_status["running"]:
        return False
    
    try:
        # 构建命令行参数
        cmd = ["bash", "./launch/run_server.sh"]
        
        # 添加配置参数
        cmd.append(f"--backbones={config['backbones']}")
        cmd.append(f"--dynamic={1 if config['dynamic'] else 0}")
        cmd.append(f"--server_port={config['server_port']}")
        cmd.append(f"--exp={config['exp']}")
        cmd.append(f"--ref_loss={config['ref_loss']}")
        cmd.append(f"--tx_power={config['tx_power']}")
        cmd.append(f"--fading_mean={config['fading_mean']}")
        cmd.append(f"--fading_var={config['fading_var']}")
        
        # 启动进程
        process = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True
        )
        
        system_status["ns3_process"] = process
        system_status["running"] = True
        system_status["error_message"] = None
        
        return True
    except Exception as e:
        system_status["error_message"] = f"启动NS-3仿真失败: {str(e)}"
        return False

# 停止NS-3仿真
def stop_ns3_simulation() -> bool:
    """
    停止NS-3仿真
    
    Returns:
        bool: 停止是否成功
    """
    global system_status
    
    if not system_status["running"]:
        return True
    
    try:
        if system_status["ns3_process"]:
            system_status["ns3_process"].terminate()
            try:
                # 等待进程终止，最多等待5秒
                system_status["ns3_process"].wait(timeout=5)
            except subprocess.TimeoutExpired:
                # 如果超时，则强制终止
                system_status["ns3_process"].kill()
            
            system_status["ns3_process"] = None
        
        system_status["running"] = False
        system_status["error_message"] = None
        
        return True
    except Exception as e:
        system_status["error_message"] = f"停止NS-3仿真失败: {str(e)}"
        return False

# Flask 路由

@app.route('/api/status', methods=['GET'])
def get_status():
    """获取系统状态"""
    return jsonify({
        "running": system_status["running"],
        "ros2_initialized": system_status["ros2_initialized"],
        "error_message": system_status["error_message"]
    })

@app.route('/api/robots', methods=['GET'])
def get_robots():
    """获取机器人状态信息"""
    # 此处应该实现检测所有机器人的逻辑
    # 示例返回固定数据
    robots = [
        {"id": 1, "status": "online", "position": [0, 0, 120]},
        {"id": 2, "status": "online", "position": [50, 50, 120]},
        {"id": 3, "status": "offline", "position": [100, 100, 0]}
    ]
    return jsonify(robots)

@app.route('/api/network/config', methods=['GET', 'POST'])
def network_config():
    """获取或设置网络配置"""
    global current_network_config
    
    if request.method == 'GET':
        return jsonify(current_network_config)
    
    elif request.method == 'POST':
        try:
            new_config = request.json
            
            # 更新配置
            for key, value in new_config.items():
                if key in current_network_config:
                    current_network_config[key] = value
            
            return jsonify({"success": True, "config": current_network_config})
        except Exception as e:
            return jsonify({"success": False, "error": str(e)}), 400

@app.route('/api/network/start', methods=['POST'])
def start_network():
    """启动NS-3网络仿真"""
    result = start_ns3_simulation(current_network_config)
    
    if result:
        return jsonify({"success": True})
    else:
        return jsonify({
            "success": False,
            "error": system_status["error_message"]
        }), 500

@app.route('/api/network/stop', methods=['POST'])
def stop_network():
    """停止NS-3网络仿真"""
    result = stop_ns3_simulation()
    
    if result:
        return jsonify({"success": True})
    else:
        return jsonify({
            "success": False,
            "error": system_status["error_message"]
        }), 500

@app.route('/api/network/data', methods=['GET'])
def get_network_data():
    """获取最新的网络数据"""
    # 添加时间戳
    data = {
        **latest_network_data,
        "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
    }
    return jsonify(data)

@app.route('/api/logs/ns3', methods=['GET'])
def get_ns3_logs():
    """获取NS-3日志"""
    # 此处应该实现读取NS-3日志文件的逻辑
    # 示例返回固定数据
    return jsonify({
        "logs": ["NS-3日志尚未实现", "请在实际部署中添加日志文件读取逻辑"]
    })

@app.route('/api/logs/ros', methods=['GET'])
def get_ros_logs():
    """获取ROS日志"""
    # 此处应该实现读取ROS日志文件的逻辑
    # 示例返回固定数据
    return jsonify({
        "logs": ["ROS日志尚未实现", "请在实际部署中添加日志文件读取逻辑"]
    })

# WebSocket 事件处理

@socketio.on('connect')
def handle_connect():
    """处理客户端连接"""
    print("客户端已连接")

@socketio.on('disconnect')
def handle_disconnect():
    """处理客户端断开连接"""
    print("客户端已断开连接")

@socketio.on('request_network_data')
def handle_request_network_data():
    """处理客户端请求网络数据"""
    data = {
        **latest_network_data,
        "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
    }
    emit('network_data', data)

# 主函数
if __name__ == '__main__':
    # 启动ROS2节点线程
    ros2_thread = threading.Thread(target=run_ros2_node)
    ros2_thread.daemon = True
    ros2_thread.start()
    
    # 启动Flask服务
    # 注意：生产环境应该使用更可靠的WSGI服务器
    socketio.run(app, host='0.0.0.0', port=5000, debug=True)
